use crate::dynf::code::{CodeObject,Instruction,Constant,BinaryOperator,Label,UnaryOperator};
use crate::dynf::{parser_program};
use crate::dynf::ast::{Program,Statement,Expression,self};

#[derive(Clone, Copy)]
struct CompileContext {
    in_loop: bool,
    func: FunctionContext,
}

#[derive(Clone, Copy, PartialEq)]
enum FunctionContext {
    NoFunction,
    Function
}

impl CompileContext {
    fn in_func(self) -> bool {
        match self.func {
            FunctionContext::NoFunction => false,
            _ => true,
        }
    }
}

pub struct Compiler {
    source_path: Option<String>,
    output_stack: Vec<CodeObject>,
    ctx: CompileContext,
    next_label: usize
}

impl Compiler {
    fn new() -> Self {
        Compiler {
            source_path: None,
            output_stack: vec![],
            ctx:CompileContext {
                in_loop:false,
                func: FunctionContext::NoFunction
            },
            next_label: 0
        }
    }

    fn emit(&mut self, instruction: Instruction) {
        self.current_output().emit(instruction);
    }

    fn new_label(&mut self) -> Label {
        let l = Label::new(self.next_label);
        self.next_label += 1;        
        l
    }

    fn set_label(&mut self, label: Label) {
        self.current_output().set_label(label)
    }

    fn current_output(&mut self) -> &mut CodeObject {
        self.output_stack.last_mut().expect("No OutputStream on stack")
    }

    fn push_output(&mut self, code: CodeObject) {
        self.output_stack.push(code);
    }

    fn pop_code_object(&mut self) -> CodeObject {
        self.output_stack.pop().unwrap()
    }

    fn push_new_code_object(&mut self, obj_name: &str) {
        let new_code_object = CodeObject::new(&obj_name);
        self.push_output(new_code_object);
    }

    fn compile_program(&mut self,program:&Program) {
        self.compile_statements(&program.statements)
    }

    fn compile_statements(&mut self, statements: &[Statement]) {
        for statement in statements {
            self.compile_statement(statement)
        }
    }

    fn compile_statement(&mut self, statement: &Statement) {
        println!("Compiling {:?}", statement);
        match statement {
            Statement::Let {sym_name,value} => {
                self.compile_expression(value);
                self.emit(Instruction::StoreName(sym_name.to_owned()));
            },
            Statement::FunctionDef {name,args,body} => {
                self.compile_function_def(name, args, body);
            },
            Statement::Return{value } => {
                if !self.ctx.in_func() {
                    panic!("InvalidReturn");
                }
                match value {
                    Some(val) => {
                        self.compile_expression(val);
                    },
                    None => {
                        self.emit(Instruction::LoadConst(Constant::None));
                    }
                }
                self.emit(Instruction::ReturnValue);
            },
            Statement::If {test,block,orelse} => {
                let end_label = self.new_label();
                match orelse {
                    None => {
                        self.compile_jump_if(test,false, end_label);
                        self.compile_statements(block);
                    },
                    Some(statements) => {
                        let else_label = self.new_label();
                        self.compile_jump_if(test, false, else_label);
                        self.compile_statements(block);
                        self.emit(Instruction::Jump(end_label));
                        // else:
                        self.set_label(else_label);
                        self.compile_statements(statements);
                    }
                }
                self.set_label(end_label);
            },
            Statement::While {test,body} => self.compile_while(test, body),
            Statement::Expression(e) => self.compile_expression(e),
            Statement::Break => {
                if !self.ctx.in_loop {
                   panic!("InvalidBreak");
                }
                self.emit(Instruction::Break);
            },
            Statement::Continue => {
                if !self.ctx.in_loop {
                    panic!("InvalidContinue");
                 }
                 self.emit(Instruction::Continue);
            },
            Statement::Assign {name,value} => {
                self.compile_expression(value);
                self.compile_store(name);
            }
            _ => ()
        }
    }

    fn compile_store(&mut self, target: &ast::Expression) {
        match  target {
            ast::Expression::Identifier(name) => {
                self.emit(Instruction::StoreName(name.to_owned()));
            },
            ast::Expression::IndexOp{idx,value} => {
                self.compile_expression(value);
                self.compile_expression(idx);
                self.emit(Instruction::StoreList);
            },
            _ => ()
        }
    }

    fn compile_while(&mut self,test:&Expression,body:&Vec<Statement>) {
        let start_label = self.new_label();
        let end_label = self.new_label();
        let break_label = self.new_label();
        self.emit(Instruction::SetupLoop { start: start_label, end: end_label});
        self.set_label(start_label);
        self.compile_jump_if(test, false, break_label);

        let was_in_loop = self.ctx.in_loop;
        self.ctx.in_loop = true;
        self.compile_statements(body);
        self.ctx.in_loop = was_in_loop;
        self.emit(Instruction::Jump(start_label));
        self.set_label(break_label);
        self.emit(Instruction::PopBlock);
        self.set_label(end_label);
    }

    fn compile_jump_if(&mut self,test:&Expression, condition: bool,target_label: Label) {
        let mut other = || {
            self.compile_expression(test);
            if condition {
                self.emit(Instruction::JumpIfTrue(target_label));
            } else {
                self.emit(Instruction::JumpIfFalse(target_label));
            }
        };
        match test {
            Expression::BoolOp {a,op,b} => {
                match op {
                    ast::BoolOperator::And => {
                        if condition {
                            let end_label = self.new_label();
                            self.compile_jump_if(a, false, end_label);
                            self.compile_jump_if(b, true, target_label);
                            self.set_label(end_label);
                        } else {
                            self.compile_jump_if(a, false, target_label);
                            self.compile_jump_if(b, false, target_label);
                        }
                    },
                    ast::BoolOperator::Or => {
                        if condition {
                            self.compile_jump_if(a, true, target_label);
                            self.compile_jump_if(b, true, target_label);
                        } else {
                            let end_label = self.new_label();
                            self.compile_jump_if(a, true, end_label);
                            self.compile_jump_if(b, false, target_label);
                            self.set_label(end_label);
                        }
                    },
                    _ => other()
                }
            },
            Expression::UnaryOp {a, op} => {
                match op {
                    ast::UnOpType::Not => {
                        self.compile_jump_if(a, !condition, target_label);
                    }
                }
            }
            _ =>  other()
        }
    }

    fn compile_function_def(&mut self,name:&String,args:&Vec<String>,body:&Vec<Statement>) {
        let prev_ctx = self.ctx;
        self.ctx = CompileContext {
            in_loop: false,
            func:FunctionContext::Function
        };
        self.enter_function(name, args);
        self.compile_statements(body);

        match body.last() {
            Some(Statement::Return { .. }) => (),
            _ => {
                self.emit(Instruction::LoadConst(Constant::None));
                self.emit(Instruction::ReturnValue);
            }
        }
        let mut code = self.pop_code_object();
        self.emit(Instruction::LoadConst(Constant::Code{code} ) );
        self.emit(Instruction::LoadConst(Constant::String {value:name.to_owned()} )  );
        self.emit(Instruction::MakeFunction);
        self.emit(Instruction::StoreName(name.to_owned()));
        self.ctx = prev_ctx;
    }

    fn enter_function(&mut self,name:&String,args:&Vec<String>) {
        let mut code_object = CodeObject::new(name);
        code_object.arg_names = args.iter().map(|a|a.clone()).collect();
        self.push_output(code_object);
    }


    fn compile_expression(&mut self, expression: &Expression) {
        match expression {
            Expression::True => {
                self.emit(Instruction::LoadConst(Constant::Boolean { value:true } ));
            },
            Expression::False => {
                self.emit(Instruction::LoadConst(Constant::Boolean { value:false } ));
            },
            Expression::Number(value) => {
                let const_value = match value {
                    ast::Number::Integer(int_num) => {
                        Constant::Integer { value: *int_num}
                    },
                    ast::Number::Float(float_num) => {
                        Constant::Float { value: float_num.into_inner() }
                    }
                };
                self.emit(Instruction::LoadConst(const_value));
            },
            Expression::Binop{a,op,b} => {
                self.compile_expression(a);
                self.compile_expression(b);
                self.compile_op(op)
            },
            Expression::UnaryOp {a,op} => {
                self.compile_expression(a);
                let iop = match op {
                    ast::UnOpType::Not => {
                        UnaryOperator::Not
                    }
                };
                self.emit(Instruction::UnaryOperation(iop));
            },
            Expression::BoolOp {a,op,b} => {
                self.compile_expression(a);
                self.compile_expression(b);
                self.emit(Instruction::Boolperation(op.clone()));
            },
            Expression::Identifier(name) => {
                self.load_name(name.as_str());
            },
            Expression::Call {args,function} => {
                self.compile_call(args, function)
            },
            Expression::List(lst) => {
                for elem in lst {
                    self.compile_expression(elem);
                }
                self.emit(Instruction::BuildList(lst.len()));
            },
            Expression::IndexOp {idx,value} => {
                self.compile_expression(value);
                self.compile_expression(idx);
                self.emit(Instruction::ListIndex);
            },
            _ => ()
        }
    }

    fn compile_call(&mut self,args:&Vec<Expression>,function:&Box<Expression>) {
        self.compile_expression(function);
        for elem in args {
            self.compile_expression(elem);
        }
        self.emit(Instruction::CallFunction(args.len()) );
    }

    fn compile_op(&mut self, op: &ast::Operator) {
        let i = match op {
            ast::Operator::Add => BinaryOperator::Add,
            ast::Operator::Sub => BinaryOperator::Subtract,
            ast::Operator::Mul => BinaryOperator::Multiply,
            ast::Operator::Div => BinaryOperator::Divide,

        };
        self.emit(Instruction::BinaryOperation(i));
    }

    fn load_name(&mut self, name: &str) {
        self.emit(Instruction::LoadName(name.to_owned()));
    }
}

pub fn compile_file(path:&str) -> CodeObject {
    let main_source = std::fs::read_to_string(path).unwrap();
    println!("{}",main_source);
    let program = parser_program(&main_source);
    //dbg!(&program);
    let mut compiler = Compiler::new();
    compiler.source_path = Some(path.to_owned());
    compiler.push_new_code_object("main");
    compiler.compile_program(&program);
    let code = compiler.pop_code_object();
    
    code
}

#[test]
fn test_compiler() {
    compile_file("main.dynf");
   
}